#!/usr/bin/env python3
# Copyright (c) 2019-2020 Windy Albert


from ctypes import BigEndianStructure
import keyTree

from addrBTC import *
from addrETH import *
from addrDOGE import *
from addrFIL import *
from addrTRX import *
from addrBNB import *

import os
import json
import sys
import getpass
import time

# m / purpose' / coin_type' / account' / change / address_index
#                                        0:externel 
#                                        1:change(internel)
# https://github.com/bitcoin/bips/blob/master/bip-0044.mediawiki#change

def usage():
    print("Make by Windy Albert version: 2.0.2")
    print("Usage:")
    print("    ./run.py view   # View the keytree by [saved] mnemonic")
    print("    ./run.py new    # Generate keytree by a [New] random mnemonic")
    print("    ./run.py clear  # Will delete all key and mnemonic [Warning]")
    print("    ./run.py lock   # Lock all cold key and mnemonic")
    print("    ./run.py unlock # Unlock all cold key and mnemonic")

def write_file(name,data):
    file = open(name,'w')
    file.flush()
    file.write(data)
    file.close()


def main():
    addrConstuct={}
    addrConstuct["0\'/0\'"] = addrBTC_p2pkh
    addrConstuct["0\'/1\'"] = addrBTC_p2wpkh
    addrConstuct["0\'/2\'"] = addrBTC_p2wpkh_p2sh
    addrConstuct["0\'/3\'"] = addrBTC_p2wsh
    addrConstuct["0\'/4\'"] = addrBTC_p2wsh_p2sh

    addrConstuct["60\'/0\'"] = addrETH
    addrConstuct["3\'/0\'"] = addrDOGE
    addrConstuct["461\'/0\'"] = addrFIL
    addrConstuct["195\'/0\'"] = addrTRX
    addrConstuct["714\'/0\'"] = addrBNB

    if len(sys.argv) != 2 :
        usage()
        return

    if not os.path.exists("secret"):
        os.mkdir("secret")    

    mfile = "secret/mnemonic.key"
    pfile = "path.conf"

    if sys.argv[1] == "new" :
        if os.path.exists(mfile):
            print("There is a mnemonic.key exist, please remove it and try again !!")
            return

        mnemonic = getRandomMemonic(random.getrandbits(160)) # for 15 word
        file = open(mfile,'w')
        file.write(json.dumps(mnemonic))
        file.close()

        print("Create mnemonic.key successfully !")
        return

    elif sys.argv[1] == "view" :
        if not os.path.exists(mfile):
            print("There is not a mnemonic.key exist, run < ./run.py new > firstly !!")
            return

        if not os.path.exists(pfile):
            print("There is not a path.conf exist, please create it firstly !!")
            return
        
        file = open(mfile,'r')
        data = file.read()
        file.close()
        mnemonic = json.loads(data)

        file1 = open(pfile,'r')
        data = file1.read()
        ptree = json.loads(data)
        file1.close()

        mnemonic1 = getRandomMemonic( int("0x"+mnemonic["rand"],0) )

        if mnemonic["mnemonic"] != mnemonic1["mnemonic"]:
            print("The mnemonic.key is [broken], fix it with the rand element ?[N/y]")
            cmd = input("")
            if cmd != "y" and cmd != "Y":
                return
            
            file = open(mfile,'w')
            file.flush()
            data = file.write( json.dumps(mnemonic1) )
            file.close()
            mnemonic = mnemonic1

        password = getpass.getpass("Input password please  : ")
        password2 = getpass.getpass("Input password again   : ")
        if password != password2 :
            print("Password is incorrect !")
            return

    elif sys.argv[1] == "clear" :
        print("Are you sure to delete all the seceret? Type in \"delete\" to continue: ")
        cmd = input("")
        if cmd == "delete":
            os.system("rm -rf Waddress* WIF secret")
            return
        else :
            print("Cancel ...")
            return

    elif sys.argv[1] == "lock" :
        os.chmod("secret",0o000)
        os.chmod("WIF/COLD",0o000)
        return

    elif sys.argv[1] == "unlock" :
        os.chmod("secret",0o770)
        os.chmod("WIF/COLD",0o770)
        return

    else :
        usage()
        return

    tree = keyTree(mnemonic["mnemonic"])

    print(COLOR_RED + "═"*136 + COLOR_NONE)
    xprv_m = tree.getMasterKey(password.encode('utf-8'))    # to b'abc'


    # write result to file
    file = open("addr.list",'w')
    file.flush()
    if not os.path.exists("WIF"):
        os.mkdir("WIF")

    # if not os.path.exists("PUB"):
    #     os.mkdir("PUB")

    address = "########################### Address List ###########################\n\n"

    purposelist = ptree["child"]
    for purpose in purposelist:
        coinlist = purpose["child"]
        purpose_key = tree.getChildpriKey(xprv_m,purpose["path"])
        for coin in coinlist:
            print_line("Coin : "+coin["coin"],COLOR_RED)
            accountlist = coin["child"]
            coin_key = tree.getChildpriKey(purpose_key,coin["path"])
            for account in accountlist:
                changelist=account["child"]
                account_key = tree.getChildpriKey(coin_key,account["path"])
            
                for change in changelist:
                    wiflist=change["child"]
                    change_prikey = tree.getChildpriKey(account_key,change["path"])
                    change_pubkey,pubkey = tree.getChildpubKey(account_key,change["path"])
                    # change_pubkey["child"] = []

                    walletObj = addrConstuct[coin["path"]+"/"+account["path"]](change_prikey)
                    for wif in wiflist:
                        wallet = walletObj.getWallet(wif["path"])
                        # change_pubkey["child"].append({"path":wif["path"]})
                        
                        if change["path"] in "0":
                            dir_wif = "WIF/COLD/"+coin["coin"]+"_"+account["account"] # ???????
                        else :
                            dir_wif = "WIF/HOT/"+coin["coin"]+"_"+account["account"]  # ???????????????????????

                        if not os.path.exists(dir_wif):
                            os.makedirs(dir_wif)

                        write_file(dir_wif+"/"+coin["coin"]+"_" + account["account"]+"_"+change["path"] +"_"+wif["path"]+".key",json.dumps(wallet))
                        address += '{:23s}'.format(coin["coin"]+"_"+account["account"]+"_"+change["path"] +"_"+wif["path"])
                        address += (": "+wallet["addr"]+"\n")

                        print('{:23s}'.format(wallet["path"])+": ")
                        name ='{:23s}'.format( coin["coin"]+"_"+account["account"] )
                        print( name + ": "+COLOR_YELLOW+wallet["wif"]+COLOR_NONE)
                        name ='{:23s}'.format( coin["coin"]+"_"+ account["account"] +"_addr " + wif["path"] )
                        print( name+": "+COLOR_GREEN+wallet["addr"]+COLOR_NONE)
                    address += "\n"

                    if "COLD" in dir_wif:
                        write_file(dir_wif +"/"+coin["coin"]+"_"+account["account"]+"_"+change["path"]+".pub",json.dumps(change_pubkey))
                    # write_file("PUB/"+coin["coin"]+"-"+account["account"]+"_"+change["path"]+".pub",json.dumps(change_pubkey))

    address += "####################################################################\n"
    file.write(address)
    file.close()

    sha256 = hashlib.new('sha256',address.encode('utf-8')).hexdigest()
    os.rename("addr.list","Waddress"+ time.strftime("_%y%m%02d_", time.localtime()) +sha256[0:6])



if __name__ == '__main__':
    main()